from __future__ import absolute_import
from __future__ import division
import sys
import os
import argparse
import yaml
import logging
from src import const
from src.segmentation import SegRegion
from src.google_tile import download_tiles_all_count
from src.geo_func import way2tiles
import numpy as np

sys.path.append('./src')
from src.inference import MyModel


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--conf', help="conf file", default="conf/server.yml", type=argparse.FileType('r'))
    parser.add_argument('-n', '--lat-n', type=int)
    parser.add_argument('-m', '--lng-n', type=int)
    parser.add_argument('-r', '--region', type=str, default='shanghai_test')
    parser.add_argument('--china-convert', type=bool, default=True)
    parser.add_argument('--dir', default='data/tiles')
    parser.add_argument('--save-file', default='data/pic_feature.txt')
    parser.add_argument('--num-classes', default=6, type=int)
    return parser.parse_args()


def gen_feature(region, model, args):
    way = region.to_way()
    tiles = way2tiles(way, args.china_convert)
    downloaded_num = download_tiles_all_count(tiles, args.dir)
    all_num = len(tiles)
    if downloaded_num / all_num < 0.8:
        logging.warning('Pic number is not enough for this region!')
        return None
    else:
        n_file = 0
        feature = np.array([0.0] * args.num_classes)
        for tile in tiles:
            filename = '%d_%d_%d.jpg' % (tile[0], tile[1], tile[2])
            filename = os.path.join(args.dir, filename)
            if os.path.exists(filename):
                one_feature, _ = model.detect_one_img(filename)
                feature += one_feature
                n_file += 1
        return (feature / n_file)


def main(args):
    conf = yaml.load(args.conf)
    model = MyModel(conf, ext='jpg')
    region = getattr(const, args.region.upper() + '_REGION')
    seg_region = SegRegion(region.left_up, region.right_down, args.lat_n, args.lng_n)
    with open(conf['label_file']) as f:
        print(f.read())
    save_file = open(args.save_file, 'w')
    for line_regions in seg_region.map:
        for one_region in line_regions:
            feature = gen_feature(one_region, model, args)
            if feature is not None:
                save_file.write('%f %f %f %f' % (one_region.left_up.lat, one_region.left_up.lng,
                                                 one_region.right_down.lat, one_region.right_down.lng))
                for x in feature:
                    save_file.write(' %f' % x)
                logging.info(feature)
                save_file.write('\n')
    save_file.close()
    with open(conf['label_file']) as f:
        print(f.read())



if __name__ == '__main__':
    args = parse_args()
    logging.basicConfig(level=logging.INFO)
    print(main(args))
